#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# @Time    : 2020/7/1 下午7:29
# @Author  : fugang_le
# @Software: PyCharm

# -*- encoding: utf-8 -*-
'''
@Author  :   wangjian
'''

from annoy import AnnoyIndex
import os
# import pymongo
import traceback
from tqdm import tqdm
import time
import json
import pickle

start_time = time.time()
with open('word2id.pkl', mode='rb') as fr:
    word2id = pickle.load(fr)

with open('id2word.pkl', mode='rb') as fr:
    id2word = pickle.load(fr)
print("load pkl time: {}ms".format(str(round((time.time() - start_time) * 1000, 2))))


class annoy_model:
    def __init__(self):
        self._load()

    def _load(self):
        cur_dir = os.path.dirname(os.path.abspath(__file__))
        model_file = os.path.join(cur_dir, 'word2vec.ann')

        self.model = AnnoyIndex(200, 'angular')
        self.model.load(model_file, prefault=True)
        # self.id_model = annoy_id_model()
        self.word2id, self.id2word = word2id, id2word
        print(len(self.word2id))
        print(len(self.id2word))

    def get_most_similar(self, word, topn=10):
        try:
            word_id = self.word2id[word]
            if word_id is None:
                return None
            else:
                start_time = time.time()
                raw_res = self.model.get_nns_by_item(i=word_id, n=topn, include_distances=True)
                print('word_id: ', word_id)
                print("get_nns_by_item time: {}ms".format(str(round((time.time() - start_time) * 1000, 2))))
                res_ids, res_scores = raw_res
                # id2word = self.id2word[res_ids]
                res_words = [self.id2word[i] for i in res_ids]
                # sqrt(2(1-cos(u,v)))
                res_scores = [self.to_cosin_distance(i) for i in res_scores]
                return (res_words, res_scores)
        except:
            print(traceback.format_exc())
            return None

    def to_cosin_distance(self, euclidean_distance):
        # sqrt(2(1-cos(u,v)))
        cosin_d = 1 - euclidean_distance ** 2 / 2
        return cosin_d  # 0.5*(1+cosin_d)


if __name__ == '__main__':
    model = annoy_model()
    res = model.get_most_similar('你好')
    print(res)
    a = ['产品说明书', '产品手册', '如何报销', '差旅申请', '费用报销', '费用查询', '差旅报销', '产品文档', '报销政策', '盲目']
    for i in range(len(a)):
        res = model.get_most_similar(a[i])
        print(res)

    while 1:
        word = input(">>")
        start_time = time.time()
        res = model.get_most_similar(word)
        print("time: {}ms".format(str(round((time.time() - start_time) * 1000, 2))))
        print(res)